import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

def im2col(input_data, kernel_size, stride=1, padding=0):
    batch_size, input_height, input_width, channels = input_data.shape
    out_height = (input_height + 2 * padding - kernel_size) // stride + 1
    out_width  = (input_width  + 2 * padding - kernel_size) // stride + 1
    output_data = np.zeros((batch_size, out_height * out_width, kernel_size * kernel_size * channels))

    for b in range(batch_size):
        for oh in range(out_height):
            for ow in range(out_width):
                for c in range(channels):
                    for kh in range(kernel_size):
                        for kw in range(kernel_size):
                            ih = oh * stride - padding + kh
                            iw = ow * stride - padding + kw
                            if 0 <= ih < input_height and 0 <= iw < input_width:
                                output_data[b, oh * out_width + ow, c * kernel_size * kernel_size + kh * kernel_size + kw] = input_data[b, ih, iw, c]
                            else:
                                output_data[b, oh * out_width + ow, c * kernel_size * kernel_size + kh * kernel_size + kw] = 0
    return output_data.reshape(batch_size, out_height * out_width, -1)

def convolution_with_im2col(input_data, weight, bias=None, stride=1, padding=0):
    batch_size, input_height, input_width, in_channels = input_data.shape
    out_channels, _, kernel_height, kernel_width = weight.shape

    im2col_matrix = im2col(input_data, kernel_height, stride, padding)
    print(f"im2col_matrix: {im2col_matrix}")

    # Reshape weight for matrix multiplication
    weight_matrix = weight.reshape(out_channels, -1)
    print(f"weight_matrix: {weight_matrix}")

    # Perform convolution as matrix multiplication
    out_height = (input_height + 2 * padding - kernel_height) // stride + 1
    out_width = (input_width + 2 * padding - kernel_width) // stride + 1
    
    output = np.zeros((batch_size, out_height * out_width, out_channels))
    for b in range(batch_size):
        output[b] = im2col_matrix[b] @ weight_matrix.T
    
    # Add bias if provided
    if bias is not None:
        output += bias
    
    # Reshape to final output format
    return output.reshape(batch_size, out_height, out_width, out_channels)


if __name__ == "__main__":
    batch_size, height, width, in_channels = 1, 5, 5, 2
    out_channels = 3
    kernel_size, stride, padding = 3, 1, 1

    input_np = np.arange(batch_size * height * width * in_channels).reshape(batch_size, height, width, in_channels).astype(np.float32)
    weight_np = np.arange(out_channels* in_channels * kernel_size * kernel_size).reshape(out_channels, in_channels, kernel_size, kernel_size).astype(np.float32)
    bias_np = np.zeros(out_channels).astype(np.float32)

    output_im2col = convolution_with_im2col(input_np, weight_np, bias_np, stride, padding)
    print(f"Output from im2col: {output_im2col}")

    input_torch = torch.from_numpy(input_np.transpose(0, 3, 1, 2))  # NHWC -> NCHW
    weight_torch = torch.from_numpy(weight_np).float()
    bias_torch = torch.from_numpy(bias_np)

    with torch.no_grad():
        output_torch = F.conv2d(input_torch, weight_torch, bias_torch, stride, padding)
        output_torch = output_torch.transpose(1, 3).transpose(1, 2).numpy()  # NCHW -> NHWC

    max_diff = np.max(np.abs(output_im2col - output_torch))
    mean_diff = np.mean(np.abs(output_im2col - output_torch))
    
    print(f"Output shape: {output_im2col.shape}")
    print(f"Max difference: {max_diff:.6f}")
    print(f"Mean difference: {mean_diff:.6f}")
    print(f"Results match: {np.allclose(output_im2col, output_torch, rtol=1e-5, atol=1e-6)}")